from __future__ import print_function
import torch
import torch.nn as nn
import torch.utils.data
import torch.nn.functional as F
from models_acv.submodule import *
import math

# 特征提取，返回1/4原尺寸分组相关性特征，1/16和1/8原尺寸特征List
class feature_extraction(nn.Module):
    def __init__(self):
        super(feature_extraction, self).__init__()

        self.inplanes = 32 # 残差层的输入通道
        self.firstconv = nn.Sequential(
            convbn(3, 32, 3, 2, 1, 1), # 2d卷积+批归一化，输入3通道，输出32通道，尺寸/2
            nn.ReLU(inplace=True),
            convbn(32, 32, 3, 1, 1, 1), # 2d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True),
            convbn(32, 32, 3, 1, 1, 1), # 2d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True)
            )
        self.layer1 = self._make_layer(BasicBlock, 32, 3, 1, 1, 1) # 构造3个残差层，无下采样块，输入32通道，输出32通道，尺寸不变
        self.layer2 = self._make_layer(BasicBlock, 64, 16, 2, 1, 1) # 构造16个残差层，有下采样块，输入32通道，输出64通道，尺寸/2
        self.layer3 = self._make_layer(BasicBlock, 128, 3, 1, 1, 1) # 构造3个残差层，有下采样块，输入64通道，输出128通道，尺寸不变
        self.layer4 = self._make_layer(BasicBlock, 128, 3, 1, 1, 2) # 构造3个残差层，无下采样块，输入128通道，输出128通道，尺寸不变

        self.feature_CFG_16 = nn.Sequential(
            BasicConv(64, 64, kernel_size=3, stride=2, padding=1), # 基本卷积块，2d卷积，通道不变，尺寸/2
            BasicConv(64, 128, kernel_size=3, stride=2, padding=1) # 基本卷积块，2d卷积，输入64通道，输出128通道，尺寸/2
            )

        self.feature_CFG_8 = nn.Sequential(BasicConv(128, 64, kernel_size=3, stride=2, padding=1)) # 基本卷积块，2d卷积，输入128通道，输出64通道，尺寸/2

    # 构造层
    def _make_layer(self, block, planes, blocks, stride, pad, dilation):
        downsample = None # 对输入下采样，保证下采样残差块的残差连接尺寸一致
        if stride != 1 or self.inplanes != planes * block.expansion: # 步长!=1或输出通道*1!=inplanes
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), # 2d卷积
                nn.BatchNorm2d(planes * block.expansion) # 批归一化
                )

        layers = [] # 用于顺序块的层列表
        layers.append(block(self.inplanes, planes, stride, downsample, pad, dilation)) # 有下采样块
        self.inplanes = planes * block.expansion # inplanes赋值为输出通道*1
        for i in range(1, blocks): # 块重复blocks-1次
            layers.append(block(self.inplanes, planes, 1, None, pad, dilation)) # 通道不变，尺寸不变

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.firstconv(x) # 前三层下采样，2d卷积+批归一化，输入3通道，输出32通道，1/2原尺寸
        x = self.layer1(x) # 3个残差层，无下采样块，输入32通道，输出32通道，尺寸不变
        l1 = self.layer2(x) # 一元特征l1，16个残差层，有下采样块，输入32通道，输出64通道，1/4原尺寸
        l2 = self.layer3(l1) # 一元特征l2(之后使用扩张卷积获得大感受野)，3个残差层，有下采样块，输入64通道，输出128通道，尺寸不变
        l3 = self.layer4(l2) # 一元特征l3(之后使用扩张卷积获得大感受野)，3个残差层，无下采样块，输入128通道，输出128通道，尺寸不变
        gwc_feature = torch.cat((l1, l2, l3), dim=1) # 分组相关性特征，在通道维度连接l1、l2、l3，输出(B,320,H/4,W/4)

        feat_16 = self.feature_CFG_16(l1) # 2d卷积下采样，输入64通道，输出128通道，1/16原尺寸
        feat_8 = self.feature_CFG_8(l2) # 2d卷积下采样，输入128通道，输出64通道，1/8原尺寸
        return {"gwc_feature": gwc_feature, "features": [feat_16, feat_8]} # 返回字典，1/4原尺寸分组相关性特征(B,320,H/4,W/4)，1/16和1/8原尺寸特征List[(B,128,H/16,W/16),(B,64,H/8,W/8)]

# 上下文与几何信息融合CGF模块
class Context_Geometry_Fusion(nn.Module):
    def __init__(self, cv_chan, im_chan): # 成本体积通道，特征通道，输出张量形状与输入的成本体积一致
        super(Context_Geometry_Fusion, self).__init__()

        self.feat = nn.Sequential(
            BasicConv(im_chan, im_chan//2, kernel_size=1, stride=1, padding=0), # 基本卷积块，2d卷积，通道/2，尺寸不变
            nn.Conv2d(im_chan//2, cv_chan, 1) # 2d卷积，改变通道，尺寸不变
            )

        self.att = nn.Sequential(
            BasicConv(cv_chan, cv_chan, is_3d=True, bn=True, relu=True, kernel_size=(1,5,5), padding=(0,2,2), stride=1, dilation=1), # 基本卷积块，3d卷积，通道不变，尺寸不变
            nn.Conv3d(cv_chan, cv_chan, kernel_size=1, stride=1, padding=0, bias=False) # 3d卷积，通道不变，尺寸不变
            )

        self.agg = BasicConv(cv_chan, cv_chan, is_3d=True, bn=True, relu=True, kernel_size=(1,5,5), padding=(0,2,2), stride=1, dilation=1) # 基本卷积块，3d卷积，通道不变，尺寸不变

    def forward(self, cv, feat): # 输入成本体积(几何信息)，特征(上下文信息)
        '''
        '''
        feat = self.feat(feat).unsqueeze(2) # 处理特征，输出形状与成本体积一致
        att = self.att(feat+cv) # 空间注意力张量(编码了要注意的信息)，成本体积和特征相加，5*5卷积
        cv = torch.sigmoid(att)*feat + cv # 加性注意力，空间注意力权重*特征(上下文信息)，融合成本体积(几何信息)，*表示对应元素相乘(哈达玛积)
        cv = self.agg(cv) # 5*5卷积

        return cv # 返回张量形状与输入的成本体积一致

# 沙漏网络，用于正则化多级补丁匹配体积和成本聚合
class hourglass(nn.Module):
    def __init__(self, in_channels):
        super(hourglass, self).__init__()

        self.conv1 = nn.Sequential(
            convbn_3d(in_channels, in_channels * 2, 3, 2, 1), # 3d卷积+批归一化，输出通道*2，尺寸/2
            nn.ReLU(inplace=True)
            )

        self.conv2 = nn.Sequential(
            convbn_3d(in_channels * 2, in_channels * 2, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True)
            )

        self.conv3 = nn.Sequential(
            convbn_3d(in_channels * 2, in_channels * 4, 3, 2, 1), # 3d卷积+批归一化，输出通道*2，尺寸/2
            nn.ReLU(inplace=True)
            )

        self.conv4 = nn.Sequential(
            convbn_3d(in_channels * 4, in_channels * 4, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True)
            )

        self.attention_block = attention_block(channels_3d=in_channels * 4, num_heads=16, block=(4, 4, 4)) # 自注意力模块，输入通道(不变)，头数量，补丁大小(在D,H,W维度上划分网格)

        self.conv5 = nn.Sequential(
            nn.ConvTranspose3d(in_channels * 4, in_channels * 2, 3, padding=1, output_padding=1, stride=2, bias=False), # 3d转置卷积，通道/2，尺寸*2
            nn.BatchNorm3d(in_channels * 2) # 3d批归一化
            )

        self.conv6 = nn.Sequential(
            nn.ConvTranspose3d(in_channels * 2, in_channels, 3, padding=1, output_padding=1, stride=2, bias=False), # 3d转置卷积，通道/2，尺寸*2
            nn.BatchNorm3d(in_channels) # 3d批归一化
            )

        self.redir1 = convbn_3d(in_channels, in_channels, kernel_size=1, stride=1, pad=0) # 3d卷积+批归一化，通道不变，尺寸不变
        self.redir2 = convbn_3d(in_channels * 2, in_channels * 2, kernel_size=1, stride=1, pad=0) # 3d卷积+批归一化，通道不变，尺寸不变

        self.CGF_16 = Context_Geometry_Fusion(in_channels*4, 128) # 上下文与几何信息融合CGF模块，输入成本体积通道，特征通道，输出张量形状与输入的成本体积一致
        self.CGF_8 = Context_Geometry_Fusion(in_channels*2, 64) # 上下文与几何信息融合CGF模块
        # self.feature_fusion_4 = Context_Geometry_Fusion(in_channels, 32)

    def forward(self, x, features): # 输入多级补丁匹配体积(B,32,D/4,H/4,W/4)，1/16和1/8原尺寸左特征List[(B,128,H/16,W/16),(B,64,H/8,W/8)]
        conv1 = self.conv1(x) # 3d卷积下采样，输入32通道，输出64通道，尺寸/2，(B,64,D/8,H/8,W/8)
        conv2 = self.conv2(conv1) # 3d卷积，通道不变，尺寸不变

        conv3 = self.conv3(conv2) # 3d卷积下采样，输入64通道，输出128通道，尺寸/2，(B,128,D/16,H/16,W/16)
        conv4 = self.conv4(conv3) # 3d卷积，通道不变，尺寸不变
        conv4 = self.attention_block(conv4) # 自注意力模块，输出(B,128,D/16,H/16,W/16)
        conv4 = self.CGF_16(conv4, features[0]) # 上下文与几何信息融合CGF模块，输入成本体积(几何信息)(B,128,D/16,H/16,W/16)，1/16原尺寸左特征(上下文信息)(B,128,H/16,W/16)，输出(B,128,D/16,H/16,W/16)

        conv5 = F.relu(self.conv5(conv4) + self.redir2(conv2), inplace=True) # 3d转置卷积上采样conv4，输入128通道，输出64通道，尺寸*2，(B,64,D/8,H/8,W/8)，残差连接conv2
        conv5 = self.CGF_8(conv5, features[1]) # 上下文与几何信息融合CGF模块，输入成本体积(几何信息)(B,64,D/8,H/8,W/8)，1/8原尺寸左特征(上下文信息)(B,64,H/8,W/8)，输出(B,64,D/8,H/8,W/8)

        conv6 = F.relu(self.conv6(conv5) + self.redir1(x), inplace=True) # 3d转置卷积上采样conv5，输入64通道，输出32通道，尺寸*2，(B,32,D/4,H/4,W/4)，残差连接x

        return conv6 # 返回(B,32,D/4,H/4,W/4)

class ACVNet(nn.Module):
    def __init__(self, maxdisp, attn_weights_only=False, freeze_attn_weights=False): # 最大视差，是否只训练注意力权重(默认False)，是否冻结注意力权重(默认False)
        super(ACVNet, self).__init__()
        self.maxdisp = maxdisp
        self.attn_weights_only = attn_weights_only # 是否只训练注意力权重(默认False)
        self.freeze_attn_weights = freeze_attn_weights # 是否冻结注意力权重(默认False)
        self.num_groups = 40
        self.concat_channels = 32
        self.feature_extraction = feature_extraction() # 特征提取
        self.concatconv = nn.Sequential(
            convbn(320, 128, 3, 1, 1, 1), # 2d卷积+批归一化，输入320通道，输出128通道，尺寸不变
            nn.ReLU(inplace=True),
            nn.Conv2d(128, self.concat_channels, kernel_size=1, padding=0, stride=1, bias=False) # 2d卷积，输入128通道，输出32通道，尺寸不变
            )

        self.patch = nn.Conv3d(40, 40, kernel_size=(1,3,3), stride=1, dilation=1, groups=40, padding=(0,1,1), bias=False) # 3d分组卷积，通道不变，尺寸不变，分40组
        self.patch_l1 = nn.Conv3d(8, 8, kernel_size=(1,3,3), stride=1, dilation=1, groups=8, padding=(0,1,1), bias=False) # 3d分组卷积，通道不变，尺寸不变，分8组
        self.patch_l2 = nn.Conv3d(16, 16, kernel_size=(1,3,3), stride=1, dilation=2, groups=16, padding=(0,2,2), bias=False) # 3d分组扩张卷积，通道不变，尺寸不变，分16组，扩张系数2(1*5*5)
        self.patch_l3 = nn.Conv3d(16, 16, kernel_size=(1,3,3), stride=1, dilation=3, groups=16, padding=(0,3,3), bias=False) # 3d分组扩张卷积，通道不变，尺寸不变，分16组，扩张系数3(1*7*7)

        self.dres1_att_ = nn.Sequential( # 正则化多级补丁匹配体积
            convbn_3d(40, 32, 3, 1, 1), # 3d卷积+批归一化，输入40通道，输出32通道，尺寸不变
            nn.ReLU(inplace=True),
            convbn_3d(32, 32, 3, 1, 1) # 3d卷积+批归一化，通道不变，尺寸不变
            )
        self.dres2_att_ = hourglass(32) # 正则化多级补丁匹配体积，沙漏网络，输入32通道

        self.classif_att_ = nn.Sequential(
            convbn_3d(32, 32, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False) # 3d卷积，输入32通道，输出1通道，尺寸不变
            )

        self.dres0 = nn.Sequential( # 成本聚合，正则化注意力串联体积
            convbn_3d(self.concat_channels * 2, 32, 3, 1, 1), # 3d卷积+批归一化，输入64通道，输出32通道，尺寸不变
            nn.ReLU(inplace=True),
            convbn_3d(32, 32, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True)
            )
        self.dres1 = nn.Sequential( # 成本聚合，正则化注意力串联体积
            convbn_3d(32, 32, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True),
            convbn_3d(32, 32, 3, 1, 1) # 3d卷积+批归一化，通道不变，尺寸不变
            )
        self.dres2 = hourglass(32) # 成本聚合，正则化注意力串联体积，沙漏网络，输入32通道
        self.dres3 = hourglass(32) # 成本聚合，正则化注意力串联体积，沙漏网络，输入32通道

        self.classif0 = nn.Sequential(
            convbn_3d(32, 32, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False) # 3d卷积，输入32通道，输出1通道，尺寸不变
            )

        self.classif1 = nn.Sequential(
            convbn_3d(32, 32, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False) # 3d卷积，输入32通道，输出1通道，尺寸不变
            )

        self.classif2 = nn.Sequential(
            convbn_3d(32, 32, 3, 1, 1), # 3d卷积+批归一化，通道不变，尺寸不变
            nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False) # 3d卷积，输入32通道，输出1通道，尺寸不变
            )

        # 权重初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.Conv3d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm3d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.bias.data.zero_()

    def forward(self, left, right):
        if self.freeze_attn_weights: # 冻结注意力权重
            with torch.no_grad():
                features_left = self.feature_extraction(left)
                features_right = self.feature_extraction(right)

                gwc_volume = build_gwc_volume(features_left["gwc_feature"], features_right["gwc_feature"], self.maxdisp // 4, self.num_groups)
                gwc_volume = self.patch(gwc_volume)
                patch_l1 = self.patch_l1(gwc_volume[:, :8])
                patch_l2 = self.patch_l2(gwc_volume[:, 8:24])
                patch_l3 = self.patch_l3(gwc_volume[:, 24:40])
                patch_volume = torch.cat((patch_l1, patch_l2, patch_l3), dim=1)

                cost_attention = self.dres1_att_(patch_volume)
                cost_attention = self.dres2_att_(cost_attention, features_left["features"])

                att_weights = self.classif_att_(cost_attention)
        else: # 不冻结注意力权重，默认运行，生成注意力权重(用于过滤初始串联体积)
            features_left = self.feature_extraction(left) # 特征提取，返回1/4原尺寸分组相关性特征(B,320,H/4,W/4)，1/16和1/8原尺寸特征List[(B,128,H/16,W/16),(B,64,H/8,W/8)]
            features_right = self.feature_extraction(right)

            gwc_volume = build_gwc_volume(features_left["gwc_feature"], features_right["gwc_feature"], self.maxdisp // 4, self.num_groups) # 构建分组相关性体积，输入1/4原尺寸左右分组相关性特征(320通道)，最大视差/4，组数(40)，输出(B,40,D/4,H/4,W/4)
            gwc_volume = self.patch(gwc_volume) # 分组相关性体积，3d分组卷积，通道不变，尺寸不变，分40组，输出(B,40,D/4,H/4,W/4)
            patch_l1 = self.patch_l1(gwc_volume[:, :8]) # 补丁特征l1，3d分组卷积，通道不变，尺寸不变，分8组，输入分组相关性体积中来自l1的部分(前8通道)，输出(B,8,D/4,H/4,W/4)
            patch_l2 = self.patch_l2(gwc_volume[:, 8:24]) # 扩张大感受野补丁特征l2，3d分组扩张卷积，通道不变，尺寸不变，分16组，扩张系数2(1*5*5)，输入分组相关性体积中来自l2的部分(8到24通道)，输出(B,16,D/4,H/4,W/4)
            patch_l3 = self.patch_l3(gwc_volume[:, 24:40]) # 扩张大感受野补丁特征l3，3d分组扩张卷积，通道不变，尺寸不变，分16组，扩张系数3(1*7*7)，输入分组相关性体积中来自l3的部分(24到40通道)，输出(B,16,D/4,H/4,W/4)
            patch_volume = torch.cat((patch_l1, patch_l2, patch_l3), dim=1) # 多级补丁匹配体积，在通道维度连接上述3个不同感受野的补丁特征，输出(B,40,D/4,H/4,W/4)

            cost_attention = self.dres1_att_(patch_volume) # 正则化多级补丁匹配体积，两个3d卷积，输入40通道，输出32通道，尺寸不变，(B,32,D/4,H/4,W/4)
            cost_attention = self.dres2_att_(cost_attention, features_left["features"]) # 正则化多级补丁匹配体积，沙漏网络，输入多级补丁匹配体积(B,32,D/4,H/4,W/4)，1/16和1/8原尺寸左特征List[(B,128,H/16,W/16),(B,64,H/8,W/8)]，输出(B,32,D/4,H/4,W/4)

            att_weights = self.classif_att_(cost_attention) # 注意力权重，3d卷积，输入32通道，输出1通道，尺寸不变，(B,1,D/4,H/4,W/4)

        if not self.attn_weights_only: # 不只训练注意力权重，默认运行
            concat_feature_left = self.concatconv(features_left["gwc_feature"]) # 串联特征，1/4原尺寸分组相关性特征(B,320,H/4,W/4)，两次2d卷积，输入320通道，输出32通道，尺寸不变，(B,32,H/4,W/4)
            concat_feature_right = self.concatconv(features_right["gwc_feature"])
            concat_volume = build_concat_volume(concat_feature_left, concat_feature_right, self.maxdisp // 4) # 构建初始串联体积，输入1/4原尺寸左右串联特征(32通道)，最大视差/4，输出(B,64,D/4,H/4,W/4)

            ac_volume = F.softmax(att_weights, dim=2) * concat_volume ### ac_volume = att_weights * concat_volume，注意力过滤，注意力串联体积=注意力权重*初始串联体积，输出(B,64,D/4,H/4,W/4)

            cost0 = self.dres0(ac_volume) # 成本聚合，正则化注意力串联体积，两个3d卷积，输入64通道，输出32通道，尺寸不变，(B,32,D/4,H/4,W/4)
            cost0 = self.dres1(cost0) + cost0 # 成本聚合，正则化注意力串联体积，两个3d卷积，通道不变，尺寸不变，残差连接，输出(B,32,D/4,H/4,W/4)
            out1 = self.dres2(cost0, features_left["features"]) # 成本聚合，正则化注意力串联体积，沙漏网络，输入注意力串联体积(B,32,D/4,H/4,W/4)，1/16和1/8原尺寸左特征List[(B,128,H/16,W/16),(B,64,H/8,W/8)]，输出(B,32,D/4,H/4,W/4)
            out2 = self.dres3(out1, features_left["features"]) # 成本聚合，正则化注意力串联体积，沙漏网络，同上

        if self.training: # 训练模式
            if not self.freeze_attn_weights: # 不冻结注意力权重，默认运行
                # 注意力权重回归视差，用于监督得到更准确的权重
                cost_attention = F.interpolate(att_weights, [self.maxdisp, left.size()[2], left.size()[3]], mode='trilinear') # 全尺寸注意力权重，三线性插值，输出(B,1,D,H,W)
                cost_attention = torch.squeeze(cost_attention, 1) # 删除大小为1的维度，输出(B,D,H,W)
                pred_attention = F.softmax(cost_attention, dim=1) # 全尺寸注意力概率体积，用于soft argmin，在视差维度做softmax，(B,D,H,W)
                pred_attention = disparity_regression(pred_attention, self.maxdisp) # 回归用于监督注意力权重的视差，输入全尺寸注意力概率体积，最大视差，输出(B,H,W)

            if not self.attn_weights_only: # 不只训练注意力权重，默认运行
                # 不同正则化程度的注意力串联体积
                cost0 = self.classif0(cost0) # 3d卷积聚合的注意力串联体积，输出(B,1,D/4,H/4,W/4)
                cost1 = self.classif1(out1) # 第一次沙漏网络聚合的注意力串联体积，输出(B,1,D/4,H/4,W/4)
                cost2 = self.classif2(out2) # 第二次沙漏网络聚合的注意力串联体积，输出(B,1,D/4,H/4,W/4)

                # 3d卷积聚合的注意力串联体积
                cost0 = F.interpolate(cost0, [self.maxdisp, left.size()[2], left.size()[3]], mode='trilinear') # 全尺寸注意力串联体积，三线性插值，输出(B,1,D,H,W)
                cost0 = torch.squeeze(cost0, 1) # 删除大小为1的维度，输出(B,D,H,W)
                pred0 = F.softmax(cost0, dim=1) # 全尺寸注意力串联概率体积，用于soft argmin，在视差维度做softmax，(B,D,H,W)
                pred0 = disparity_regression(pred0, self.maxdisp) # 回归视差0，输入全尺寸注意力串联概率体积，最大视差，输出(B,H,W)

                # 第一次沙漏网络聚合的注意力串联体积
                cost1 = F.interpolate(cost1, [self.maxdisp, left.size()[2], left.size()[3]], mode='trilinear') # 全尺寸注意力串联体积，三线性插值，输出(B,1,D,H,W)
                cost1 = torch.squeeze(cost1, 1) # 删除大小为1的维度，输出(B,D,H,W)
                pred1 = F.softmax(cost1, dim=1) # 全尺寸注意力串联概率体积，用于soft argmin，在视差维度做softmax，(B,D,H,W)
                pred1 = disparity_regression(pred1, self.maxdisp) # 回归视差1，输入全尺寸注意力串联概率体积，最大视差，输出(B,H,W)

                # 第二次沙漏网络聚合的注意力串联体积
                cost2 = F.interpolate(cost2, [self.maxdisp, left.size()[2], left.size()[3]], mode='trilinear') # 全尺寸注意力串联体积，三线性插值，输出(B,1,D,H,W)
                cost2 = torch.squeeze(cost2, 1) # 删除大小为1的维度，输出(B,D,H,W)
                pred2 = F.softmax(cost2, dim=1) # 全尺寸注意力串联概率体积，用于soft argmin，在视差维度做softmax，(B,D,H,W)
                pred2 = disparity_regression(pred2, self.maxdisp) # 回归视差2，输入全尺寸注意力串联概率体积，最大视差，输出(B,H,W)

                if self.freeze_attn_weights: # 冻结注意力权重
                    return [pred0, pred1, pred2]

                return [pred_attention, pred0, pred1, pred2] # 默认运行，返回[用于监督注意力权重的视差(B,H,W),回归视差0(B,H,W),回归视差1(B,H,W),回归视差2(B,H,W)]，后三个视差是不同成本体积聚合程度的回归结果

            return [pred_attention]
        else: # 验证模式
            if self.attn_weights_only: # 只训练注意力权重
                cost_attention = F.interpolate(att_weights, [self.maxdisp, left.size()[2], left.size()[3]], mode='trilinear')
                cost_attention = torch.squeeze(cost_attention, 1)
                pred_attention = F.softmax(cost_attention, dim=1)
                pred_attention = disparity_regression(pred_attention, self.maxdisp)

                return [pred_attention]

            cost2 = self.classif2(out2) # 第二次沙漏网络聚合的注意力串联体积，输出(B,1,D/4,H/4,W/4)
            cost2 = F.interpolate(cost2, [self.maxdisp, left.size()[2], left.size()[3]], mode='trilinear') # 全尺寸注意力串联体积，三线性插值，输出(B,1,D,H,W)
            cost2 = torch.squeeze(cost2, 1) # 删除大小为1的维度，输出(B,D,H,W)
            pred2 = F.softmax(cost2, dim=1) # 全尺寸注意力串联概率体积，用于soft argmin，在视差维度做softmax，(B,D,H,W)
            pred2 = disparity_regression(pred2, self.maxdisp) # 回归视差2，输入全尺寸注意力串联概率体积，最大视差，输出(B,H,W)

            return [pred2] # 默认运行，返回最深聚合网络的回归结果，[回归视差2(B,H,W)]

def acv(d):
    return ACVNet(d)
